import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { DataStore } from './type'

const dataStore = create<DataStore>()(
    persist(
        (set) => ({
            data: {
                original_document_id: '',
                indexing_technique: 'high_quality',
                // 处理规则
                process_rule: {
                    mode: 'automatic',
                    // 自定义规则
                    rules: {
                        pre_processing_rules: [
                            {
                                // 替换连续空格、换行符、制表符
                                id: 'remove_extra_spaces',
                                enabled: false
                            },
                            {
                                // 删除 URL、电子邮件地址
                                id: 'remove_urls_emails',
                                enabled: true
                            }
                        ],
                        //  父块 分段规则
                        segmentation: {
                            separator: '\\n2\\n',
                            max_tokens: 500
                        },
                        // 召回方式  full-doc-全文召回   paragraph-段落召回
                        parent_mode: 'paragraph',
                        // 子分段规则
                        subchunk_segmentation: {
                            separator: '\\n',
                            max_tokens: 200
                        }
                    }
                },
                // 检索模式
                retrieval_model: {
                    // 默认混合检索
                    search_method: 'hybrid_search',
                    // 默认开启rerank
                    reranking_enable: true,
                    // 模型配置
                    reranking_model: {
                        // 模型供应商
                        reranking_provider_name: 'gte-rerank',
                        // 模型名称
                        reranking_model_name: 'langgenius/tongyi/tongyi'
                    },
                    // 召回条数
                    top_k: 5,
                    // 是否开启召回分数限制
                    score_threshold_enabled: false,
                    // 召回分数限制
                    score_threshold: 0.5
                },
                // 模型名称
                embedding_model: 'text-embedding-v3',
                // 模型供应商
                embedding_model_provider: 'langgenius/tongyi/tongyi'
            },
            RAG: {
                // 名称
                // 名称
                name: '',
                // base64编码头像
                avatar: '',
                // 说明
                description: '',
                // 嵌入模型的名称
                embedding_model: 'text-embedding-v3',
                // 谁可以访问要创建的数据  "me"：（默认）只有您可以管理数据集。 "team"：所有团队成员都可以管理数据集。
                permission: 'me',
                // 创建的数据集的分块方法
                chunk_method: 'laws',
                parser_config: {
                    layout_recognize: 'Plain Text',
                    auto_keywords: 0,
                    auto_questions: 0,
                    graphrag: {
                        entity_types: ["organization", "person", "geo", "event", "category"],
                        method: 'light',
                        use_graphrag: true
                    },
                    html4excel: false,
                    delimiter: "\n!?。；！？",
                    task_page_size: 12,
                    raptor: {
                        max_cluster: 64,
                        max_token: 256,
                        prompt: '请总结以下段落。 小心数字，不要编造。 段落如下：\n      {cluster_content}\n以上就是你需要总结的内容。',
                        use_raptor: true,
                        random_seed: 0,
                        threshold: 0.1,
                        knowledge_graph: true
                    },
                }
            },
            setdata: (path, value) => {
                set((state) => {
                    const newState = { ...state };
                    const paths = path.split('.');
                    let current = newState.data;
                    for (let i = 0; i < paths.length - 1; i++) {
                        current = current[paths[i]];
                    }
                    current[paths[paths.length - 1]] = value;
                    return newState;
                });
            }
        }),
        {
            name: 'data-store', // 持久化存储的名称
            partialize: (state) => ({

            }),
        }
    )
);

export default dataStore;


// function adjustParserConfig(chunkMethod: string, parserConfig: any) {
//     const defaultConfig = {
//         chunk_token_count: 512,
//         layout_recognize: true,
//         html4excel: false,
//         delimiter: "\n!?。；！？",
//         task_page_size: 12,
//         raptor: { use_raptor: false },
//     };

//     switch (chunkMethod) {
//         case "naive":
//             return defaultConfig;
//         case "manual":
//         case "qa":
//         case "paper":
//         case "book":
//         case "laws":
//         case "presentation":
//             return {
//                 raptor: { use_raptor: false },
//             };
//         case "table":
//         case "picture":
//         case "one":
//         case "email":
//             return {};
//         case "knowledge_graph":
//             return {
//                 chunk_token_count: 512,
//                 delimiter: "\n!?。；！？",
//                 entity_types: ["organization", "person", "location", "event", "time"],
//             };
//         default:
//             return {};
//     }
// }